// Copyright 2021-present StarRocks, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <string>
#include <vector>

#include "common/status.h"
#include "fs/fs.h" // FileInfo
#include "gen_cpp/data.pb.h"
#include "gen_cpp/lake_types.pb.h"
#include "storage/tablet_schema.h"

namespace starrocks {
class Chunk;
class Column;
class TabletSchema;
class ThreadPool;

namespace lake {

class TabletManager;

enum WriterType : int { kHorizontal = 0, kVertical = 1 };

// Basic interface for tablet writers.
class TabletWriter {
public:
    explicit TabletWriter(TabletManager* tablet_mgr, int64_t tablet_id, std::shared_ptr<const TabletSchema> schema,
                          int64_t txn_id, ThreadPool* flush_pool = nullptr)
            : _tablet_mgr(tablet_mgr),
              _tablet_id(tablet_id),
              _schema(std::move(schema)),
              _txn_id(txn_id),
              _flush_pool(flush_pool) {}

    virtual ~TabletWriter() = default;

    TabletManager* tablet_manager() const { return _tablet_mgr; }

    int64_t tablet_id() const { return _tablet_id; }

    int64_t txn_id() const { return _txn_id; }

    // Return a list of file info generated by this writer.
    // Note: The path of `FileInfo` is only base filename, no directory component.
    //
    // PREREQUISITES: the writer has successfully `finish()`ed but not yet `close()`ed.
    std::vector<FileInfo> files() const { return _files; }

    // The sum of all segment file sizes, in bytes.
    int64_t data_size() const { return _data_size; }

    // The total number of rows have been written.
    int64_t num_rows() const { return _num_rows; }

    // This method is called immediately before any elements are processed, it
    // should contain the writer's initialization logic.
    virtual Status open() = 0;

    // Writes elements from the specified chunk to this rowset.
    //
    // It's guaranteed that the elements in each chunk are arranged in ascending
    // order, and the elements among all chunks written before `flush()` are also
    // arranged in ascending order.
    //
    // For horizontal writer.
    virtual Status write(const Chunk& data, SegmentPB* segment = nullptr) = 0;

    // Writes partial columns data to this rowset.
    //
    // It's guaranteed that the elements in each chunk are arranged in ascending order by keys,
    // and multiple segments will be automatically generated according to `max_rows_per_segment`.
    //
    // For vertical writer.
    virtual Status write_columns(const Chunk& data, const std::vector<uint32_t>& column_indexes, bool is_key) = 0;

    // Write del file to this rowset. PK table only
    virtual Status flush_del_file(const Column& deletes) = 0;

    // Flushes this writer and forces any buffered bytes to be written out to segment files.
    // There is no order guarantee between the data written before a `flush()`
    // and the data written after it.
    virtual Status flush(SegmentPB* segment = nullptr) = 0;

    // Flushes partial columns data when current columns are written finished.
    //
    // For vertical writer.
    virtual Status flush_columns() = 0;

    // This method is called at the end of data processing.
    virtual Status finish(SegmentPB* segment = nullptr) = 0;

    // This method is called at the very end of the operator's life, both in
    // the case of a successful completion of the operation, and in the case
    // of a failure and canceling.
    virtual void close() = 0;

    //  return the rowset txn meta generated by tablet writer
    virtual RowsetTxnMetaPB* rowset_txn_meta() = 0;

    // allow to set custom tablet schema for writer, used in partial update
    void set_tablet_schema(TabletSchemaCSPtr schema) { _schema = std::move(schema); }

protected:
    TabletManager* _tablet_mgr;
    int64_t _tablet_id;
    TabletSchemaCSPtr _schema;
    int64_t _txn_id;
    ThreadPool* _flush_pool;
    std::vector<FileInfo> _files;
    int64_t _num_rows = 0;
    int64_t _data_size = 0;
    uint32_t _seg_id = 0;
    bool _finished = false;
};

} // namespace lake
} // namespace starrocks
